import subprocess
import os
import datetime
import argparse


def take_screenshot(local_dir='.'):
    """
    使用ADB命令在连接的安卓设备上截图，
    将截图传输到本地目录，并可选地在设备上删除截图。

    :param local_dir: 保存截图的本地目录（默认：当前目录）
    """
    try:
        # 生成带有时间戳的唯一文件名以避免覆盖
        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        device_filename = f"/sdcard/screenshot_{timestamp}.png"
        local_filename = os.path.join(local_dir, f"screenshot_{timestamp}.png")

        # 第一步：在设备上截图
        subprocess.run(['adb', 'shell', 'screencap', '-p', device_filename], check=True)
        print(f"设备上已截图：{device_filename}")

        # 第二步：将截图传输到本地目录
        subprocess.run(['adb', 'pull', device_filename, local_filename], check=True)
        print(f"截图已传输到：{local_filename}")

        # 第三步：从设备上删除截图（可选，以节省空间）
        subprocess.run(['adb', 'shell', 'rm', device_filename], check=True)
        print(f"设备上的截图已删除：{device_filename}")

    except subprocess.CalledProcessError as e:
        print(f"执行ADB命令出错：{e}")
        print("请确保ADB已安装，设备已连接，且已启用USB调试。")
    except FileNotFoundError:
        print("未找到ADB。请确保ADB已安装并添加到系统PATH中。")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="ADB截图工具")
    parser.add_argument("--dir", default=".", help="保存截图的本地目录（默认：当前目录）")
    args = parser.parse_args()

    while 1:
        take_screenshot(local_dir=args.dir)
        input("如需再次截图请回车确认")